"""
海龟交易策略(Zipline实现)
策略逻辑：
1. 使用20日突破作为入场信号
2. 使用10日低点作为止损点
3. 使用ATR计算仓位大小
"""

from zipline.api import order_target_percent, record, symbol, get_datetime
from zipline.finance import commission, slippage
import numpy as np

def initialize(context):
    # 策略参数
    context.entry_period = 20
    context.exit_period = 10
    context.atr_period = 14
    context.risk_pct = 0.01
    context.asset = symbol('AAPL')
    
    # 设置交易成本
    context.set_commission(commission.PerShare(cost=0.001, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.01))
    
    # 初始化状态
    context.in_position = False
    context.atr = 0

def handle_data(context, data):
    # 获取历史数据
    prices = data.history(context.asset, ['high', 'low', 'close'], 
                         max(context.entry_period, context.atr_period) + 1, '1d')
    
    # 计算突破通道
    entry_high = prices['high'][-context.entry_period:].max()
    exit_low = prices['low'][-context.exit_period:].min()
    
    # 计算ATR
    high = prices['high']
    low = prices['low']
    close = prices['close']
    
    tr = np.maximum(high - low, 
                   np.maximum(np.abs(high - close.shift(1)), 
                             np.abs(low - close.shift(1))))
    atr = tr.rolling(context.atr_period).mean().iloc[-1]
    
    current_price = data.current(context.asset, 'price')
    
    # 交易逻辑
    if not context.in_position and current_price > entry_high:
        # 计算仓位大小
        position_size = (context.risk_pct * context.portfolio.portfolio_value) / atr
        order_target_percent(context.asset, position_size)
        context.in_position = True
        context.atr = atr
        
    elif context.in_position and current_price < exit_low:
        order_target_percent(context.asset, 0)
        context.in_position = False
        
    # 记录状态
    record(price=current_price, 
          entry_high=entry_high, 
          exit_low=exit_low, 
          atr=atr, 
          position=context.portfolio.positions[context.asset].amount * current_price)